"""
YOLOv8 训练脚本 (与预处理脚本完全兼容)
使用全局配置进行模型训练
"""
import os
# 必须在所有import之前设置！
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'  # 解决OpenMP冲突
os.environ['OMP_NUM_THREADS'] = '1'  # 限制OpenMP线程数
import yaml
from pathlib import Path
from ultralytics import YOLO
import torch
import logging

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("logs/training.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class YOLOv8Trainer:
    def __init__(self, config_path="configs/default.yaml"):
        """
        初始化训练器
        :param config_path: 全局配置文件路径
        """
        self.config = self.load_config(config_path)
        self.setup_paths()
        self.setup_env()

    def load_config(self, config_path):
        """加载并验证配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                cfg = yaml.safe_load(f)

            # 必要参数检查
            required_keys = {
                'dataset', 'training', 'paths', 'settings'
            }
            if not required_keys.issubset(cfg.keys()):
                missing = required_keys - cfg.keys()
                raise ValueError(f"Missing required config keys: {missing}")

            return cfg
        except Exception as e:
            logger.error(f"配置文件加载失败: {str(e)}")
            raise

    def setup_paths(self):
        """处理路径配置"""
        base_dir = Path(__file__).parent

        # 确保dataset.yaml使用绝对路径
        self.dataset_yaml = str((base_dir / self.config['paths']["processed"]["dataset"]).resolve())

        # 项目输出目录不应放在data目录下，改为独立目录
        self.project_dir = str((base_dir / self.config["paths"]["runs"]["train"]).resolve())

        # 验证数据集文件存在
        if not Path(self.dataset_yaml).exists():
            raise FileNotFoundError(
                f"dataset.yaml 不存在于: {self.dataset_yaml}\n"
                f"请确保预处理脚本已正确生成数据集配置"
            )

        # 创建输出目录
        os.makedirs(self.project_dir, exist_ok=True)

    def setup_env(self):
        """设置训练环境"""
        # 设备配置
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        logger.info(f"使用设备: {self.device}")

        # 训练参数默认值
        self.train_config = {
            'model': 'yolov8n.pt',
            'epochs': 100,
            'batch': 16,
            'imgsz': 640,
            'workers': 4,
            'name': 'exp',
            'exist_ok': True,
            'hyperparameters': {
                'lr0': 0.01,
                'lrf': 0.01,
                'momentum': 0.937,
                'weight_decay': 0.0005,
                'warmup_epochs': 3
            }
        }

        # 合并用户自定义配置
        if 'training' in self.config:
            self.train_config.update(self.config['training']['base'])

    def train(self):
        """执行训练流程"""
        try:
            # 加载模型
            model = YOLO(self.train_config['model'])

            # 训练参数
            train_args = {
                'data': self.dataset_yaml,
                'project': self.project_dir,
                'device': self.device,
                **{k: v for k, v in self.train_config.items()
                   if k not in ['model', 'hyperparameters']},
                **self.train_config['hyperparameters']
            }

            logger.info("开始训练，参数配置:")
            logger.info(yaml.dump(train_args, allow_unicode=True))

            # 在训练前设置环境变量
            os.environ["ULTRALYTICS_DATASETS"] = str(Path(self.dataset_yaml).parent)

            # 启动训练
            results = model.train(**train_args)

            # 保存最终配置
            self.save_final_config()

            return results
        except Exception as e:
            logger.error(f"训练过程出错: {str(e)}")
            raise

    def save_final_config(self):
        """保存最终训练配置"""
        config_path = Path(self.project_dir) / self.train_config['name'] / "train_config.yaml"
        with open(config_path, 'w') as f:
            yaml.dump({
                'dataset_config': self.config,
                'train_config': self.train_config
            }, f, sort_keys=False)


if __name__ == "__main__":
    try:
        logger.info("=" * 50)
        logger.info("YOLOv8 训练启动")
        logger.info("=" * 50)

        trainer = YOLOv8Trainer()
        results = trainer.train()

        logger.info("\n" + "=" * 50)
        logger.info(f"训练完成! 模型保存在: {trainer.project_dir}/{trainer.train_config['name']}")
        logger.info("=" * 50)
    except Exception as e:
        logger.error(f"训练失败: {str(e)}")
        exit(1)